import akshare as ak
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from statsmodels.tsa.arima.model import ARIMA
from sklearn.metrics import mean_absolute_error, mean_squared_error
from keras.models import Sequential
from keras.layers import LSTM, Dense
from sklearn.preprocessing import MinMaxScaler
import warnings

warnings.filterwarnings("ignore")  # Ignore warning messages


# ==================== Data Fetching ====================
def get_china_stock_data(stock_code="000001", adjust="hfq"):
    """Fetch China A-share stock data (default: Ping An Bank, adjusted)"""
    try:
        df = ak.stock_zh_a_hist(symbol=stock_code, adjust=adjust)
        df['Date'] = pd.to_datetime(df['日期'])  # Changed to English
        df.set_index('Date', inplace=True)
        df = df[['收盘']].rename(columns={'收盘': 'Close'})

        # Data cleaning: fill missing values
        df = df.ffill().bfill()  # Forward then backward fill

        # Ensure no NaN values
        if df.isnull().values.any():
            raise ValueError("Data contains NaN values, please check data source")

        df.to_csv('stock_data.csv')
        print(f"Successfully downloaded {stock_code} data to stock_data.csv")
        return df
    except Exception as e:
        print(f"Failed to fetch data: {e}")
        return None


# ==================== ARIMA Model ====================
def arima_predict(data, forecast_days=7):
    # Ensure no NaN values
    data = data.dropna()

    # Set frequency (handle holiday gaps)
    data = data.asfreq('D').ffill()

    train_size = int(len(data) * 0.8)
    train, test = data[:train_size], data[train_size:]

    try:
        model = ARIMA(train, order=(5, 1, 0))
        model_fit = model.fit()

        # Forecasting
        forecast = model_fit.forecast(steps=forecast_days)
        test_forecast = model_fit.forecast(steps=len(test))

        # Ensure no NaN in predictions
        test_forecast = pd.Series(test_forecast).fillna(method='ffill').values
        forecast = pd.Series(forecast).fillna(method='ffill').values

        # Calculate metrics
        mae = mean_absolute_error(test.values, test_forecast)
        rmse = np.sqrt(mean_squared_error(test.values, test_forecast))

        # Visualization
        plt.figure(figsize=(12, 6))
        plt.plot(train.index, train, label='Training Data')
        plt.plot(test.index, test, label='Actual Price')
        plt.plot(test.index, test_forecast, label='Test Prediction')

        future_dates = pd.date_range(start=data.index[-1], periods=forecast_days + 1, freq='D')[1:]
        plt.plot(future_dates, forecast, label=f'ARIMA Forecast\nMAE: {mae:.2f}, RMSE: {rmse:.2f}', linestyle='--')

        plt.title('ARIMA Stock Price Prediction')
        plt.legend()
        plt.show()

        return forecast, mae, rmse
    except Exception as e:
        print(f"ARIMA model error: {e}")
        return None, None, None


# ==================== LSTM Model ====================
def lstm_predict(data, forecast_days=7, time_step=30):
    # Ensure no NaN values
    data = data.dropna()

    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(data)

    def create_dataset(dataset, look_back=1):
        X, Y = [], []
        for i in range(len(dataset) - look_back - 1):
            X.append(dataset[i:(i + look_back), 0])
            Y.append(dataset[i + look_back, 0])
        return np.array(X), np.array(Y)

    train_size = int(len(scaled_data) * 0.8)
    train, test = scaled_data[:train_size], scaled_data[train_size:]
    X_train, y_train = create_dataset(train, time_step)
    X_test, y_test = create_dataset(test, time_step)

    X_train = np.reshape(X_train, (X_train.shape[0], time_step, 1))
    X_test = np.reshape(X_test, (X_test.shape[0], time_step, 1))

    model = Sequential([
        LSTM(50, return_sequences=True, input_shape=(time_step, 1)),
        LSTM(50),
        Dense(1)
    ])
    model.compile(loss='mean_squared_error', optimizer='adam')
    model.fit(X_train, y_train, epochs=10, batch_size=16, verbose=1)

    test_predict = model.predict(X_test)
    test_predict = scaler.inverse_transform(test_predict)
    y_test = scaler.inverse_transform([y_test]).flatten()

    mae = mean_absolute_error(y_test, test_predict.flatten())
    rmse = np.sqrt(mean_squared_error(y_test, test_predict.flatten()))

    last_window = X_test[-1]
    future_predictions = []
    for _ in range(forecast_days):
        next_pred = model.predict(last_window.reshape(1, time_step, 1))
        future_predictions.append(next_pred[0, 0])
        last_window = np.append(last_window[1:], next_pred)

    future_predictions = scaler.inverse_transform([future_predictions])[0]

    future_dates = pd.date_range(start=data.index[-1], periods=forecast_days + 1, freq='D')[1:]

    plt.figure(figsize=(12, 6))
    plt.plot(data.index[:train_size], data.iloc[:train_size], label='Training Data')
    plt.plot(data.index[train_size + time_step:len(data) - 1],
             data.iloc[train_size + time_step:len(data) - 1], label='Actual Price')
    plt.plot(data.index[train_size + time_step:len(data) - 1],
             test_predict, label='Test Prediction')
    plt.plot(future_dates, future_predictions,
             label=f'LSTM Forecast\nMAE: {mae:.2f}, RMSE: {rmse:.2f}', linestyle='--')
    plt.title('LSTM Stock Price Prediction')
    plt.legend()
    plt.show()

    return future_predictions, mae, rmse


# ==================== Main Program ====================
if __name__ == "__main__":
    # Fetch data (example: 600519 Kweichow Moutai)
    stock_data = get_china_stock_data("600519")

    if stock_data is not None:
        # Check data quality
        print("\nData Overview:")
        print(stock_data.head())
        print(f"\nData Points: {len(stock_data)}")
        print(f"Contains NaN values: {stock_data.isnull().values.any()}")

        # ARIMA Prediction
        print("\nRunning ARIMA model...")
        arima_results, arima_mae, arima_rmse = arima_predict(stock_data)

        if arima_results is not None:
            # LSTM Prediction
            print("\nRunning LSTM model...")
            lstm_results, lstm_mae, lstm_rmse = lstm_predict(stock_data)

            # Result Comparison
            print("\n=== Model Comparison ===")
            print(f"ARIMA - MAE: {arima_mae:.4f}, RMSE: {arima_rmse:.4f}")
            print(f"LSTM  - MAE: {lstm_mae:.4f}, RMSE: {lstm_rmse:.4f}")

            # Forecast Results Table
            future_dates = pd.date_range(start=stock_data.index[-1], periods=8, freq='D')[1:]
            result_df = pd.DataFrame({
                'Date': future_dates,
                'ARIMA_Forecast': arima_results,
                'LSTM_Forecast': lstm_results
            })
            print("\nNext 7 Days Forecast:")
            print(result_df.round(2))